# -*- coding: utf-8 -*-
"""
******* 文档说明 ******

TF.Keras 训练模型 转换成 PB Serving 模型

保存成saved_model格式，自动形成以下文件目录
/save/dir/path
└─ 1
    ├── saved_model.pb
    └── variables

TF Serving Docker 服务启动命令
docker run --rm --name Vincent_TF_Serving  -p 8500:8500 -p 8501:8501 -v .\TFServing:/models \
-t tensorflow/serving:1.12.3 --model_config_file=/models/model.config --platform_config_file=/models/platform.config

# 当前项目: Cifar10-Classification
# 开发作者: vincent
# 创建平台: PyCharm Community Edition  TensorFlow 1.12.0
# 版    本: V1.0
"""
import os
import tensorflow as tf
from tensorflow import keras


# Keras h5 模型转换
def keras_model(model_name, model_version, model_path, serving_model_path=None):
    """
    :param model_name:          模型命名
    :param model_version:       模型版本字符 可 int 类型  '20190909101001'
    :param model_path:          Keras 训练好 .h5 模型 (模型结构+权重值)
    :param serving_model_path:  转换后的Serving模型保存路径，若为 None， 默认为 model_path 下路径
    :return:
    """
    tf.reset_default_graph()
    sess = tf.Session()
    keras.backend.set_session(sess)

    # ############## 导入模型
    model = keras.models.load_model(model_path)

    # 获取keras模型的输入
    input_tensor = {}
    for i, i_input in enumerate(model.inputs):
        input_tensor['input_{:d}'.format(i)] = i_input

    # 获取keras模型的输出
    output_tensor = {}
    for i, i_output in enumerate(model.outputs):
        output_tensor['output_{:d}'.format(i)] = i_output

    print(input_tensor)
    print(output_tensor)

    # # 版本时间信息，根据实际情况修改此行代码
    # model_version = (os.path.basename(model_path).split('_')[1])[:8]
    assert int(model_version), 'Error:  Serving Model Name must be int!'

    # 模型保存路径
    if serving_model_path is None:
        serving_model_path = os.path.abspath(os.path.join(os.path.dirname(__file__), model_name, model_version))
    else:
        serving_model_path = os.path.abspath(os.path.join(serving_model_path, model_name, model_version))

    # 导出模型
    tf.saved_model.simple_save(sess, export_dir=serving_model_path,
                               inputs=input_tensor, outputs=output_tensor)

    print('Serving Model {} Saved in {}'.format(model_name, serving_model_path))

if __name__ == '__main__':
    # Keras h5 模型转换
    keras_model(model_name='lpOCR',  # 模型名称
                model_version='20190921001709',  # 版本信息必须为可 Int 类型
                model_path=r'D:\Desktop\LP_project\_model\LP_model.lp-densenet2_20190921-001701-Base_Translated.h5',
                serving_model_path=r'D:\TFServing'  # pb 模型保存路径
                )
